/// File: VSLabViewlets\function-viewer3d.fs
/// 
/// Main author: 
///   Don Syme (dsyme@microsoft.com)
///
/// Adaptation to viewlet:
///   Sara Berardelli (berardel@di.unipi.it)
///   Antonio Cisternino (cisterni@di.unipi.it)
///   
/// ------------------------------------------------------------
/// Copyright (c) Microsoft Corporation.  All rights reserved.
/// 
/// The use and distribution terms for this software are 
/// contained in the file named license.txt, which can be found 
/// in the root of this distribution.
/// By using this software in any fashion, you are agreeing to 
/// be bound by the terms of this license.
///
/// You must not remove this notice, or any other, from this
/// software.
/// ------------------------------------------------------------


//----------------------------------------------------------------------------
//  FUNCTION VIEWER 3D
//  
//----------------------------------------------------------------------------
#light

module VSLab3DViewer

#I @"C:\Windows\Microsoft.NET\DirectX for Managed Code\1.0.2902.0"
#r @"Microsoft.DirectX.dll"
#r @"Microsoft.DirectX.Direct3D.dll"
#r @"Microsoft.DirectX.Direct3DX.dll"

open VSLab

//---------------------------------------------------------------
// THE LIBRARY
//---------------------------------------------------------------
open System
open System.ComponentModel
open Drawing
open Windows.Forms
open Microsoft.DirectX
open Microsoft.DirectX.Direct3D
open List

//---------------------------------------------------------------
// THE SAMPLE
//---------------------------------------------------------------
type time = float
type event<'a> = ('a -> unit) * IEvent<'a> 

type view = { mutable ypr: Matrix
              mutable focus: Vector3
              mutable zoom: float }

type moveEvent = { dx: int
                   dy: int
                   dir: int
                   key: int }

type Function3DViewlet() as x = 
  inherit Viewlet() as base

  //----------------------------------------------------------------------------
  // Configuration
  //----------------------------------------------------------------------------
  let axisColor = ref Color.White
  let bgColor = ref Color.Black
  /// function to display
  let ff = ref (fun (t:float) (x,y) -> 
                let x,y = x - 0.5,y - 0.5
                let r = sqrt (x*x + y*y)
                exp(-4.0 * r) * sin(6.0 * 3.14 * r + t) + 0.5)

  //----------------------------------------------------------------------------
  // common
  //----------------------------------------------------------------------------
  let mapR  a b f = List.init (b-a+1) (fun i -> f (a+i))
  let now () = (float Environment.TickCount / 1000.0) 
  let pi = Math.PI
  
  //----------------------------------------------------------------------------
  // vectors and operations
  //----------------------------------------------------------------------------
  let vector (x,y,z) = new Vector3(Float32.of_float x,Float32.of_float y,Float32.of_float z)    
  let vO  = vector (0.0, 0.0, 0.0)
  let vX  = vector (1.0, 0.0, 0.0)
  let vY  = vector (0.0, 1.0, 0.0)
  let vZ  = vector (0.0, 0.0, 1.0)
  let scale k u = Vector3.Scale(u,k)
  let colorPoint (c:Color) v = new CustomVertex.PositionNormalColored(v, vZ, c.ToArgb())
  let point v = colorPoint !axisColor v
  let transformAt v m = Matrix.Translation(v) * m * Matrix.Translation(-v)  

  //----------------------------------------------------------------------------
  // Color interpolation
  //----------------------------------------------------------------------------
  let colorRange x =
    let r,g,b =
      if x < 0.0 then 1.0,0.0,0.0
      elif x < 0.5 then 
        let z = 2.0 * x
        1.0 - z,z,0.0 
      elif x < 1.0 then
        let z = 2.0 * x - 1.0 
        0.0,1.0 - z,z   
      else 0.0,0.0,1.0
    let byte x = int (x * 255.0)
    Color.FromArgb(0,byte r,byte g,byte b)

  //----------------------------------------------------------------------------
  // DirectX - device
  //----------------------------------------------------------------------------
  /// presentation parameters
  let defaultPresentParams () = 
    let presentParams = new PresentParameters(Windowed = true,
                                              SwapEffect = SwapEffect.Discard,
                                              EnableAutoDepthStencil = true,
                                              AutoDepthStencilFormat = DepthFormat.D16)
    presentParams
  
  let rec checkResetThen (device: Device) presentParams f = 
    if device.CheckCooperativeLevel() then f()
    else
      let r = ref 0 
      if not (device.CheckCooperativeLevel(r)) then 
        if !r = int ResultCode.DeviceNotReset then
          device.Reset([|presentParams|])
          checkResetThen device presentParams f
  let checkVisibleThen (form : #Control) f = if form.Visible then f()
  let clearScene (color:Color) (device:Device) =
    device.Clear(ClearFlags.ZBuffer ||| ClearFlags.Target,color, 1.0f, 0)
  
  /// render function
  let doRender device presentParams form fireEvent = 
    checkResetThen device presentParams (fun () -> 
      checkVisibleThen form (fun () ->
        // Avoid inconsistent devices
        try 
          device.BeginScene()
          clearScene !bgColor device
          fireEvent()
          device.EndScene()
          device.Present()
        with _ -> ()
      )
    )

  /// initialize device
  let doInitialise(device:Device) (form: #Control) =
    device.RenderState.ZBufferEnable <- true;
    device.RenderState.Ambient <- Drawing.Color.White;
    form.Invalidate()

  /// setup lighting
  let setupLights (device:Device) =
    let mutable material = new Material(DiffuseColor = ColorValue.FromColor(Color.White),
                                        AmbientColor = ColorValue.FromColor(Color.White))
    device.Material <- material
    device.RenderState.Lighting <- true
    (device.Lights.Item(0)).Type      <- LightType.Directional
    (device.Lights.Item(0)).Diffuse   <- System.Drawing.Color.White
    (device.Lights.Item(0)).Direction <- new Vector3(0.0f,0.0f,-1.0f)
    (device.Lights.Item(0)).Enabled   <- true
    (device.Lights.Item(1)).Type      <- LightType.Point
    (device.Lights.Item(1)).Diffuse   <- System.Drawing.Color.White
    (device.Lights.Item(1)).Position  <- new Vector3(0.5f, 0.5f, -0.5f)
    (device.Lights.Item(1)).Enabled   <- true
    device.RenderState.Ambient <- System.Drawing.Color.FromArgb(0x101010)

  //----------------------------------------------------------------------------
  // DirectX - VertexBuffers
  //----------------------------------------------------------------------------
  let vertexBufferOfPoints (device:Device) (pts : CustomVertex.PositionNormalColored[]) =
    let vertexBuffer = new VertexBuffer(typeof<CustomVertex.PositionNormalColored>,
                                        pts.Length, // number pts 
                                        device,     
                                        Usage.None,
                                        CustomVertex.PositionNormalColored.Format,
                                        Pool.Managed) // Pool.Managed resources survive device loss 
    let offset = 0
    vertexBuffer.SetData(pts,offset,LockFlags.None)
    vertexBuffer
  
  //----------------------------------------------------------------------------
  // DirectX - draw lines and triangle strips
  //----------------------------------------------------------------------------
  let drawPrimitive vertexBuffer primitive nPrim (device:Device) =
    device.SetStreamSource(0, vertexBuffer, 0)
    device.VertexFormat <- CustomVertex.PositionNormalColored.Format;
    device.DrawPrimitives(primitive, 0, nPrim)
    
  let drawLineList pts device =
    let ptsA = pts |> Array.of_list
    use vb = vertexBufferOfPoints device ptsA
    drawPrimitive vb PrimitiveType.LineList (ptsA.Length/2) device

  let drawTriangeStrip pts device =
    let ptsA = pts |> Array.of_list
    use vb = vertexBufferOfPoints device ptsA
    drawPrimitive vb PrimitiveType.TriangleStrip (ptsA.Length-2) device

  //----------------------------------------------------------------------------
  // Surface - index
  //----------------------------------------------------------------------------
  let triangleRow n m =
    mapR 0 (2*n-1) (fun k ->
    let i = k / 2
    if k % 2 = 0 then (i,m) else (i,m+1))
  let triangleRows n m = mapR 0 (m-2) (fun m -> triangleRow n m)

  //----------------------------------------------------------------------------
  // setView
  //----------------------------------------------------------------------------
  let move view moveEvent =
    if moveEvent.dir = int(MouseButtons.Left) || 
      moveEvent.dir = int(Keys.Left) ||
      moveEvent.dir = int(Keys.Right) ||
      moveEvent.dir = int(Keys.Up) ||
      moveEvent.dir = int(Keys.Down) then
    
      if moveEvent.key = int(Keys.Shift) then
        view.zoom <- view.zoom * exp (float moveEvent.dy / 100.0) // Zoom 
      else
        let rx = float32 moveEvent.dx / 20.0f
        let ry = float32 moveEvent.dy / 20.0f    
        let m = Matrix.RotationYawPitchRoll(ry,0.0f,rx)    // Rotate 
        let m = transformAt (scale -0.5f (vX + vY + vZ)) m // at centre point 
        view.ypr <- m * view.ypr
  
    elif moveEvent.dir = int(MouseButtons.Right) || 
      moveEvent.dir = int(Keys.A) ||
      moveEvent.dir = int(Keys.D) ||
      moveEvent.dir = int(Keys.W) ||
      moveEvent.dir = int(Keys.S) then
    
      let dv = scale (float32 moveEvent.dx / 50.0f) vY + scale (float32 (-moveEvent.dy) / 50.0f) vZ
      view.focus <- view.focus + dv // Move focus
   
  let setView view (device:Device) =
    let eye = scale 2.0f (vX + vY + vZ) + scale (float32 view.zoom) vX
    device.Transform.View <-
      Matrix.Invert(view.ypr) * Microsoft.DirectX.Matrix.LookAtLH(eye, view.focus, vZ)
    device.Transform.Projection <-
      Microsoft.DirectX.Matrix.PerspectiveFovLH(float32 (Math.PI / 8.0),  // FOV 
                                                1.0f,                    // aspect 
                                                0.1f,                    // min depth 
                                                100.0f                   // max depth 
      )
    device.Transform.World <- Matrix.Identity
  
  //----------------------------------------------------------------------------
  // XY grid and axis box
  //----------------------------------------------------------------------------
  let planeN = 6  // number of division on XY plane grid 
  let planeVs = 
    mapR 0 planeN (fun i -> 
      let k = float32 i / float32 planeN
      [scale k vY; vX + scale k vY;      // Line k.Y  to    X + k.Y 
       scale k vX; vY + scale k vX; ]    // Line k.X  to  k.X +   Y 
    ) |> concat
  let planePts = map (colorPoint Color.Gray) planeVs
  let updateBoxPts() =
    map point
      [vO; vZ;
       vO; vX;
       vX; vX + vY;
      ] @
    map (colorPoint Color.Gray)
      [vY      ; vY      + vZ;
       vX + vY ; vX + vY + vZ;
       vZ ; vY      + vZ;
       vY + vZ ; vX + vY + vZ;     
      ]
      
  let mutable boxPts = updateBoxPts()



  let drawAxisName (device:Device) view (names:string array) =
    let localFont = new System.Drawing.Font("Times New Roman", 0.1f, FontStyle.Italic)
    use xAxisMesh = Mesh.TextFromFont(device, localFont, names.[0], 0.0f, 0.1f)
    use yAxisMesh = Mesh.TextFromFont(device, localFont, names.[1], 0.0f, 0.1f)
    use zAxisMesh = Mesh.TextFromFont(device, localFont, names.[2], 0.0f, 0.1f)  

    device.Transform.World <- Matrix.Scaling(0.1f,0.1f,0.1f) * Matrix.Translation(vX) * Matrix.RotationAxis(vX, 1.62f)
    xAxisMesh.DrawSubset(0)
    device.Transform.World <- Matrix.Identity
    device.Transform.World <- Matrix.Scaling(0.1f,0.1f,0.1f) * Matrix.Translation(-vZ) * Matrix.RotationAxis(vX, 1.62f)
    yAxisMesh.DrawSubset(0)
    device.Transform.World <- Matrix.Identity
    device.Transform.World <- Matrix.Scaling(0.1f,0.1f,0.1f) * Matrix.Translation(vY) * Matrix.RotationAxis(vX, 1.62f)
    zAxisMesh.DrawSubset(0)  
    device.Transform.World <- Matrix.Identity

  let drawAxis (device:Device) view names =
    device.RenderState.CullMode <- Cull.None
    drawLineList planePts device
    drawLineList boxPts device
    drawAxisName device view names

  //----------------------------------------------------------------------------
  // mesh and gridLines
  //----------------------------------------------------------------------------
  let mkMesh gx gy (m,n) = Array2.init n m (fun i j -> gx n i),Array2.init n m (fun i j -> gy m j)
  let meshDims mesh = 
      let X,Y = mesh 
      Array2.length1 X, Array2.length2 X
  let meshGet mesh (i,j) = 
      let X,Y = mesh 
      X.[i,j], Y.[i,j]

  let ij2k mesh (i,j) = 
      let m,n = meshDims mesh  
      i + j * n 
  let k2ij mesh k = 
      let m,n = meshDims mesh
      k % n,k / n   
  let colorPlace mesh data c (i,j) =
    let k = ij2k mesh (i,j)
    let x,y,z = data.[k]
    let x,y,z = (float32 x), (float32 y), (float32 z)
    colorPoint c (vO + scale x (vX - vO) + scale y (vY - vO) + scale z (vZ - vO))
  let blendPlace mesh data (i,j) =
    let k = ij2k mesh (i,j)
    let x,y,z = data.[k]
    let c = colorRange z
    let x,y,z = (float32 x), (float32 y), (float32 z)
    colorPoint c (vO + scale x vX + scale y vY + scale z vZ)
    
  let normi n i = float i / float (n-1)
  let grid (m,n) = mkMesh normi normi (m,n)
  let mesh = ref (grid (20,20))
  let gridLine i j  = [i,j;i+1,j;  i,j;i,j+1]
  let gridLines n m = mapR 0 (n-2) (fun i -> mapR 0 (m-2) (gridLine i)) |> concat |> concat

  let drawSurf f (device:Device) =
    let m,n = meshDims !mesh 
    let data = Array.init (n*m) (fun k ->
                                   let i,j = k2ij !mesh k
                                   let x,y = meshGet !mesh (i,j)
                                   let z = f (x,y)
                                   (x,y,z))
    let strips = triangleRows n m |> map (map (blendPlace !mesh data))
    List.iter (fun strip -> drawTriangeStrip strip device) strips;
    let lines = gridLines n m
    let lines = lines |> map (colorPlace !mesh data Color.Black)
    drawLineList lines device
    
  //----------------------------------------------------------------------------
  // state: evaluation function
  //----------------------------------------------------------------------------
  /// Z-range
  let range = ref (0.0,1.5)
  let eval t (x,y) =
    let z = !ff t (x,y)
    let a,b = !range // and scale w.r.t. range
    let res = if b - a = 0.0 then 0.0 else (z-a) / (b-a)
    res

  //----------------------------------------------------------------------------
  // change the grid
  //----------------------------------------------------------------------------
  let cgrid1 fd m = 
    let dxs = mapR 0 (m-2) (normi (m-1) >> fd)
    let sumx = fold_right (+) dxs 0.0
    let _,xs = fold_right (fun dx (sum,res) -> (sum+dx),(sum/sumx :: res)) dxs (0.0,[])
    let arr = Array.of_list (1.0 :: xs)
    fun i -> arr.[i]
  let cgrid fdensity (m,n) = mkMesh (cgrid1 fdensity) (cgrid1 fdensity) (m,n)
  (*mesh := cgrid (fun x -> (((x - 0.5) * 2.0) ** 2.0) + 0.2) (20,20)*)


  // The drawScene event fires each time we draw the scene.
  let drawSceneE : event<time> = IEvent.create()
  let timer =  new System.Timers.Timer(50.0)
  let presentParams = defaultPresentParams() 
  let mutable device : Device = null
  let axisName = [| "x"; "y"; "z" |] 
  
  let view =  
      { ypr   = Matrix.RotationYawPitchRoll(0.0f,0.0f,0.0f)
        focus = scale 0.5f (vX + vY + vZ)
        zoom  = 2.0 }

  let lastArgs = ref None

  let mutable name = "VSLab 3D function"
  member x.Name
    with set(v) = 
      name <- v
      base.Name <- v

  override x.Show() =
    base.Show()    
    device <- new Device(0, DeviceType.Hardware, base.ViewletHandle,
                         CreateFlags.SoftwareVertexProcessing, 
                         [| presentParams |])    
    x.SetStyle(ControlStyles.AllPaintingInWmPaint ||| ControlStyles.Opaque, true)
    timer.Elapsed.Add(fun e -> x.Invalidate())
    device.DeviceReset.Add(fun _ -> doInitialise device x)
    doInitialise device x
    let DrawScene = snd drawSceneE
    DrawScene.Add(fun t -> setView view device)
    DrawScene.Add(fun t -> drawAxis device view axisName)
    DrawScene.Add(fun t -> drawSurf (eval t) device)
    DrawScene.Add(fun _ -> setupLights device)
    timer.Start()
          
  member x.DrawScene = snd drawSceneE
  member x.Device = device
  member x.Timer = timer
  
  member x.XAxisName
    with get() = axisName.[0]
    and set (s:string) =
        axisName.[0] <- s
        x.Invalidate()
  
  member x.YAxisName
    with get() = axisName.[1]
    and set (s:string) =
        axisName.[1] <- s
        x.Invalidate()
  
  member x.ZAxisName
    with get() = axisName.[2]
    and set (s:string) =
        axisName.[2] <- s
        x.Invalidate()
  
  member x.Function
    with get() = !ff
    and set(f) =
        ff := f
        x.Invalidate()
  
  [<Category("Style")>]
  member x.AxisColor
    with get() = !axisColor
    and set(c) = 
      axisColor := c
      boxPts <- updateBoxPts()
      x.Invalidate()
  [<Category("Style")>]
  member x.Background
    with get() = !bgColor
    and set(c) = 
      bgColor := c
      x.Invalidate()

  
  override x.OnMouseDown args =
    lastArgs := Some args
    
  override x.OnMouseUp _ =
    lastArgs := None

  override x.OnMouseMove args =
   match !lastArgs with
      | Some last -> 
        let ev = {dx = args.X - last.X
                  dy = last.Y - args.Y
                  dir = int(args.Button)
                  key = int(Form.ModifierKeys) }
        move view ev
        lastArgs := Some args
      | None -> ()
      
  override x.IsInputKey (k: Keys) = true

  override x.OnKeyDown e =
    let hor, ver = 
        match e.KeyCode with
        | Keys.Left | Keys.A -> (-5,0)
        | Keys.Right | Keys.D -> (5,0)
        | Keys.Up | Keys.W -> (0,5)
        | Keys.Down | Keys.S -> (0,-5)
        | _ -> (0,0)

    let ev = {dx = hor
              dy = ver
              dir = int(e.KeyCode)
              key = int(e.Modifiers) }
    move view ev

  override x.OnPaintBackground(evt) = ()

  override x.OnPaint(evt:PaintEventArgs) =
    if device <> null then
      let doDrawScene,drawScene = drawSceneE 
      doRender device presentParams x (fun () -> doDrawScene (now()))
  
  override x.OnViewletVisibilityChange (v) =
    timer.Enabled <- v
    
  interface IDisposable with
    member x.Dispose () =
      x.Timer.Stop()
